#!/home/berald01/.local/bin/python

import sys
import os
import re
import subprocess
import shutil
import time
import argparse


parser = argparse.ArgumentParser(description= """

DESCRIPTION:
    Execute bwa on a fastq file by splitting it in n lines and aligning them
    splitted files in parallel.
    Output files are *.bam and *.clean.bam

    Genomes and their tags (might not be up-to-date!):
genome_dict= { '/lustre/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa': 'mm9',
               '/lustre/sblab/berald01/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa': 'mm9',
               '/lustre/sblab/berald01/reference_data/genomes/Homo_sapiens_NCBI_v36.3/hsa.fa': 'hg18'}  

EXAMPLES:
    NOTE: Use &!
    bwa_aling_split.py myfastq.fq &

    for fq in `ls mb*`
    do
    bwa_align_split.py -f $fq -g /lustre/sblab/berald01/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa --notidyup &
    done
    
TODO:
    - Start alignment after a split file has been generated. Don;t wait for the entire
      fastq to be splitted.
    - Better way of checking if a bjob has completed. Use lsf's tools for this
    - Allow the splitting of not gzipped files (skip gunzip -c)
    - Send log files to separate dir?
    
""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('-f', '--fastq',
                    type= str,
                    required= True,
                    help="""Fastq file to align. Must be gzipped with extension
.fastq.gz or .fq.gz.
                    """)

parser.add_argument('-g', '--genome', # nargs='?', default=sys.stdout,
                    type= str,
                    required= True,
                    help="""Path to genome to align to. 
                    """)

parser.add_argument('--bwaopt',
                    required= False,
                    default= '',
                    help="""Strings of options to pass to bwa.
                    """)

parser.add_argument('-o', '--outdir',
                    required= False,
                    default= '/lustre/sblab/berald01/repository/bam',
                    help="""Output dir for bam files, unfiltered. Default is
/lustre/sblab/berald01/repository/bam
                    """)
parser.add_argument('-O', '--outdirclean',
                    required= False,
                    default= '/lustre/sblab/berald01/repository/bam_clean',
                    help="""Output dir for bam files, CLEAN. Default:
/lustre/sblab/berald01/repository/bam_clean
                    """)

parser.add_argument('-w', '--whitelist',
                    required= False,
                    default= '',
                    help="""'White list': BED file of regions to include in the filtered
bam file. This file is the opposite of the black-list regions from encode. For hg18
use '/home/berald01/reference_seqs/wgEncodeHg18DukeSignalArtifactRegions.whitelist.bed'
                    """)

parser.add_argument('-m', '--mem',
                    required= False,
                    default= 5120,
                    type= int,
                    help="""Memory usage to pass to -R "rusage[mem=MEM]". Default
to 5120.
                    """)

parser.add_argument('-q', '--mapq',
                    required= False,
                    default= 15,
                    type= int,
                    help="""Minimum mapq score accepted in the clean bam file.
This int passed to 'samtools view -q int'. Default 15.
                    """)

parser.add_argument('-r', '--reads',
                    required= False,
                    default= 1000000,
                    type= int,
                    help="""Number of reads each split files will have. MEMO: This
is number of reads not number of lines. Default 1000000 (4M lines).
                    """)

parser.add_argument('-x', '--fastqext',
                    required= False,
                    default= '\..*',
                    type= str,
                    help="""String with a regex to pass to re.sub() to remove
this match from the fastq file and obtain the basename for the bam files. Default is '\..*' 
whch will strip everything fromthe left of the first dot found.
                    """)

parser.add_argument('--notidyup',
                    action= 'store_true',
                    help="""With this flag the split dir <fastq>_split will not
be removed. Useful for debugging.
                    """)


args = parser.parse_args()

# -----------------------------------------------------------------------------
FASTQGZ= args.fastq ## '/lustre/sblab/berald01/repository/fastq/test.fq.gz' ## sys.argv[1]
GENOME= args.genome ## '/lustre/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa' ## '/lustre/sblab/berald01/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa'

# FQ_DIR= '/lustre/sblab/berald01/jobs/bwa/'      ## All the files in this dir will be processed.
FASTQ_EXT= args.fastqext ## Regex to get fastq files. This match will be stripped off the name and the rest used for the output files. MEMO: regex '\.s_[1-8]_sequence\.txt\.gz$' to strip full Solexa name from suffix.
BWA_ALN_OPT= args.bwaopt
OUTPUT_DIR= args.outdir ## os.path.join(os.path.split(FQ_DIR)[0], 'bwa_aligned') will go to one level up the fastq dir and will create it.
OUTPUT_DIR_CLEAN= args.outdirclean ## Where to put bam files after having removed mapq 5 and blacklist regions

WHITELIST= args.whitelist ## Inverse of black-list. Regions to keep in clean bam output. Use '/home/berald01/reference_seqs/wgEncodeHg18DukeSignalArtifactRegions.whitelist.bed' for hg18. Use None or '' if N/A
MEM= args.mem

genome_dict= { '/lustre/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa': 'mm9',
               '/lustre/sblab/berald01/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa': 'mm9',
               '/lustre/sblab/berald01/reference_data/genomes/Homo_sapiens_NCBI_v36.3/hsa.fa': 'hg18',
               '/lustre/sblab/berald01/reference_data/genomes/Homo_sapiens_NCBI_v36.3/gap_telomere.fa': 'hg18_telomere'
               } ## Tag for output files to look like '/lustre/.../fastq/bham-470.bwa.mm9.bam'

nlines= args.reads * 4

## ----------------------------------------------------------------------------

fastq= re.sub('\.gz$', '', FASTQGZ)
splitdir= re.sub('\.fq$|\.fastq$', '', fastq) + '_split'

genome_tag= genome_dict[GENOME]
bamname= re.sub(FASTQ_EXT, '', FASTQGZ) + '.bwa.' + genome_tag + '.bam'
bamclean_name= re.sub(FASTQ_EXT, '', FASTQGZ) + '.bwa.' + genome_tag + '.clean.bam'



## Prepare split files
try:
    shutil.rmtree(splitdir)
except OSError:
    pass
os.makedirs(splitdir)

cmd= 'bsub -R "rusage[mem=1024]" -J %(fastqgz)s -oo %(fastqgz)s.split.log "gunzip -c %(fastqgz)s | split -l %(nlines)s -d - %(splitdir)s/; touch %(splitdone)s"' %{'fastqgz':FASTQGZ, 'nlines': nlines, 'splitdir':splitdir, 'splitdone': os.path.join(splitdir, 'splitdone')}
print(cmd)
p= subprocess.Popen(cmd, shell= True)

while True:
    """ Check splitting has completed by testing the presence of check-file
    TODO: Use lsf tools for this!!
    """
    time.sleep(5)
    if os.path.exists(os.path.join(splitdir, 'splitdone')):
        os.remove(os.path.join(splitdir, 'splitdone'))
        break

## Get files to align
fastqfiles= sorted(os.listdir(splitdir))

if WHITELIST is not None and WHITELIST != '':
    whitelist= '-L %s ' %(WHITELIST)
else:
    whitelist= ' '    

for fq in fastqfiles:
    """
    Prepare output file names
    """
    print(fq)
    basename= re.sub(FASTQ_EXT, '', fq)
    fq_path= os.path.join(splitdir, fq)
    sai= os.path.join(splitdir, basename + '.bwa.' + genome_tag + '.sai')
    sam= os.path.join(splitdir, basename + '.bwa.' + genome_tag + '.sam')
    bam= os.path.join(splitdir, basename + '.bwa.' + genome_tag + '.unsorted.bam')
    bam_sorted= os.path.join(splitdir, basename + '.bwa.' + genome_tag + '.bam')
    bam_sorted_clean= os.path.join(splitdir, basename + '.bwa.' + genome_tag + '.clean.bam')
    cmd_file= os.path.join(splitdir, fq + '.sh')
    """
    Compile commands: Each element of this list will be in the .sh file
    """
    cmd_list= ['bwa aln %s %s %s > %s' %(BWA_ALN_OPT, GENOME, fq_path, sai),
               'bwa samse %s %s %s > %s' %(GENOME, sai, fq_path, sam),
               'samtools view -bS %s > %s' %(sam, bam),
               'samtools sort %s %s' %(bam, re.sub('\.bam$', '', bam_sorted)),
               'rm %s; rm %s; rm %s' %(sai, sam, bam),
               'samtools view -b -q %s %s %s > %s' %(args.mapq, whitelist, bam_sorted, bam_sorted_clean)
              ]
    """
    Write commands to file(s)
    """
    shfile= open(cmd_file, 'w')
    shfile.write('#!/bin/sh' + '\n\n' + 'set -e' + '\n\n')
    for cmd in cmd_list:
        #print(cmd + '\n\n')
        shfile.write(cmd + '\n\n')
    shfile.close()
    """
    Submit jobs
    """
    bjob= '''bsub -w 'ended("%s")' -R "rusage[mem=%s]" -J %s -oo %s < %s''' %(FASTQGZ, MEM, FASTQGZ + '-' + fq, cmd_file + '.log', cmd_file)
    print(bjob)
    subprocess.Popen(bjob,shell= True)

## Join files toghether
cmd= '''bsub -w 'ended("%(jobid)s*")' -R "rusage[mem=2048]" -J %(mergeid)s -oo %(log)s "samtools merge -f %(outbam)s %(splitdir)s/*.%(genome)s.bam; samtools index %(outbam)s; touch %(mergedone)s" ''' %{'jobid':FASTQGZ, 'log': FASTQGZ + '.merge.log', 'mergeid':'merger-' + bamname ,'outbam': os.path.join(OUTPUT_DIR, bamname), 'splitdir': splitdir, 'genome': genome_tag, 'mergedone': os.path.join(splitdir, 'mergedone')}
subprocess.Popen(cmd, shell= True)
cmd= '''bsub -w 'ended("%(jobid)s*")' -R "rusage[mem=2048]" -J %(mergeid)s -oo %(log)s "samtools merge -f %(outbam)s %(splitdir)s/*.%(genome)s.clean.bam; samtools index %(outbam)s;touch %(mergedone)s" ''' %{'jobid':FASTQGZ, 'log': FASTQGZ + '.mergeclean.log', 'mergeid':'merger-' + bamclean_name, 'outbam': os.path.join(OUTPUT_DIR_CLEAN, bamclean_name), 'splitdir': splitdir, 'genome': genome_tag, 'mergedone': os.path.join(splitdir, 'mergecleandone')}
subprocess.Popen(cmd, shell= True)

## Tidy up
if args.notidyup is False:
    while True:
        """ Check mergeing has completed by testing the presence of check-file
        TODO: Use lsf tools for this!!
        """
        time.sleep(5)
        if os.path.exists(os.path.join(splitdir, 'mergedone')) & os.path.exists(os.path.join(splitdir, 'mergecleandone')):
            shutil.rmtree(splitdir)
            break
sys.exit()